use crate::key_manager::error::KeyManagerError;
use anyhow::Result;
use attestation_common::log::info;
use mockall::automock;
use once_cell::sync::Lazy;
use openssl::pkey::PKey;
use openssl::pkey::Private;
use openssl::pkey::Public;
use parking_lot::Mutex;
use std::collections::HashMap;

// factory function
pub fn create_algorithm(algorithm_str: &str) -> Result<Box<dyn KeyAlgorithm>, KeyManagerError> {
    info!("create_algorithm: {}", algorithm_str);
    let parts: Vec<&str> = algorithm_str.split_whitespace().collect();
    let (name, args) = parts
        .split_first()
        .ok_or_else(|| KeyManagerError::new("Empty algorithm string"))?;

    let registry = ALGORITHM_REGISTRY.lock();
    let ctor = registry
        .get(*name)
        .ok_or_else(|| KeyManagerError::new(format!("Algorithm not found: {}", name)))?;
    info!("create_algorithm success: {}", algorithm_str);
    Ok(ctor(args)?)
}

// Global algorithm registry
type AlgorithmConstructor = fn(&[&str]) -> Result<Box<dyn KeyAlgorithm>>;
pub(crate) static ALGORITHM_REGISTRY: Lazy<Mutex<HashMap<&'static str, AlgorithmConstructor>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

#[automock]
pub trait KeyAlgorithm {
    fn derive_public(&self, private: &Vec<u8>) -> Result<Vec<u8>, KeyManagerError>;
    fn sign(&self, private: &PKey<Private>, data: Vec<u8>) -> Result<Vec<u8>, KeyManagerError>;
    fn verify(
        &self,
        public: &PKey<Public>,
        data: Vec<u8>,
        signature: Vec<u8>,
    ) -> Result<bool, KeyManagerError>;
}

#[macro_export]
macro_rules! register_algorithm {
    ($name:expr, $ctor:expr) => {
        #[ctor::ctor]
        fn register() {
            $crate::key_manager::algorithm::factory::algorithm_factory::ALGORITHM_REGISTRY
                .lock()
                .insert($name, $ctor);
        }
    };
}

mod tests {
    #![allow(warnings)]
    use super::*;
    use mockall::{mock, predicate::eq};
    use openssl::rsa::Rsa;

    #[test]
    fn test_create_algorithm() {
        let mock_ctor = |args: &[&str]| -> Result<Box<dyn KeyAlgorithm>> {
            let mut mock_algo = MockKeyAlgorithm::new();
            mock_algo
                .expect_derive_public()
                .with(eq(vec![1, 2, 3]))
                .times(1)
                .returning(|_| Ok(vec![4, 5, 6]));

            mock_algo
                .expect_sign()
                .withf(|private, data| {
                    data == &vec![7, 8, 9]
                })
                .times(1)
                .returning(|_, _| Ok(vec![10, 11, 12]));

            mock_algo
                .expect_verify()
                .withf(|public, data, signature| {
                    // DER code for comparing public keys
                    data == &vec![13, 14, 15]
                })
                .times(1)
                .returning(|_, _, _| Ok(true));

            Ok(Box::new(mock_algo))
        };
        // Register the Mock constructor to the global registry
        ALGORITHM_REGISTRY
            .lock()
            .insert("mock_algorithm", mock_ctor);
        let algorithm = create_algorithm("mock_algorithm").unwrap();
        let public_key = algorithm.derive_public(&vec![1, 2, 3]).unwrap();
        assert_eq!(public_key, vec![4, 5, 6]);

        let private_key = Rsa::generate(2048).unwrap().private_key_to_pem().unwrap();
        let signature = algorithm
            .sign(
                &PKey::private_key_from_pem(&private_key).unwrap(),
                vec![7, 8, 9],
            )
            .unwrap();
        assert_eq!(signature, vec![10, 11, 12]);

        let public_key = PKey::private_key_from_pem(&private_key).unwrap().public_key_to_pem().unwrap();
        let result = algorithm
            .verify(
                &PKey::public_key_from_pem(&public_key).unwrap(),
                vec![13, 14, 15],
                vec![16, 17, 18],
            )
            .unwrap();
        assert!(result);
    }
}
